import logging
import fvcore.nn.weight_init as weight_init
from typing import Optional
import torch
from torch import nn, Tensor
from torch.nn import functional as F
import numpy as np

from detectron2.config import configurable
from detectron2.layers import Conv2d
from detectron2.structures import Boxes, ImageList, Instances, BitMasks

from .helper import get_aggregator

# @TRANSFORMER_DECODER_REGISTRY.register()
class DefaultInstanceRetrievalDecoder(nn.Module):
    # @configurable
    def __init__(
        self,
        input_channel=256
    ):
        super().__init__()
        
        dG = 1024
        dL = 1024
        #---- Aggregator
        # agg_arch='CosPlace'
        # agg_config={'in_dim': 256,
        #             'out_dim': 2048}
        # agg_arch='GeM'
        # agg_config={'p': 3}
        
        # agg_arch='ConvAP'
        # agg_config={'in_channels': 256,
        #             'out_channels': 512}
        self.align_size = 32
        self.out_feature_dim = 2048
        agg_arch='MixVPR'
        agg_config={'in_channels' : input_channel,
                'in_h' : self.align_size,
                'in_w' : self.align_size,
                'out_channels' : 512,
                'mix_depth' : 4,
                'mlp_ratio' : 1,
                'out_rows' : 4}
        # agg_arch='NetVLAD'
        # agg_config={'in_channels': 256,
        #             'out_channels': 2048}
        self.aggregator = get_aggregator(agg_arch,agg_config)
        self.input_proj = None
        
    # @classmethod
    def from_config(cls, cfg):
        ret = {}
        ret["num_classes"] = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES
        return ret

    ##in multi_scale_features B C H W 
    '''
    EXAMPLE
    SRC IMAGE 10*512*512
    features:res2 10, 256, 128, 128
    features:res2 10, 512, 64, 64
    features:res2 10, 1024, 32, 32
    features:res2 10, 2048, 16, 16
    mask_features :10, 256, 128, 128
    transformer_encoder_features :10, 256, 16, 16
    multi_scale_features[0]:10, 256, 16, 16
    multi_scale_features[1]:10, 256, 32, 32
    multi_scale_features[2]:10, 256, 64, 64
    
    b_id_boxes:
        int tensor:b,id,x0,y0,x1,y1 
    '''
    def forward(self,mask_features,boxes,btensor):
        if self.input_proj is not None:
            mask_features = self.input_proj(mask_features)
        bs,c,h4,w4 = mask_features.shape
        mix_feats_tensor = torch.tensor([]).to(mask_features)      
        for bnum in range(bs):
            zzf = self.inference(mask_features,boxes,btensor,bnum)
            if zzf is not None:
                mix_feats_tensor = torch.cat((mix_feats_tensor,zzf),dim=0)
        return mix_feats_tensor
    
    def inference(self,mask_features,boxes,btensor,bnum):
        b_index = torch.where(btensor == bnum)
        b_box = boxes[b_index]
        if(len(b_box) < 1):
            return None
        mask_feature = mask_features[[bnum]]
        
        boxex_features = torch.tensor([]).to(mask_feature)
        align_size = self.align_size
        
        for box in b_box:
            mask_feature_box = mask_feature[...,box[0]:box[2],box[1]:box[3]]
            mask_feature_box = F.interpolate(
                mask_feature_box, size=(align_size,align_size), mode="bilinear", align_corners=False
            )
            boxex_features = torch.cat((boxex_features,mask_feature_box),dim=0)     
        
        zzf = self.aggregator(boxex_features)
        return zzf
    
    def inference_one(self,mask_feature,boxes):
        if(len(boxes) < 1):
            return None
        if self.input_proj is not None:
            mask_feature = self.input_proj(mask_feature)
                   
        boxex_features = torch.tensor([]).to(mask_feature)
        align_size = self.align_size
        
        for box in boxes:
            mask_feature_box = mask_feature[...,box[0]:box[2],box[1]:box[3]]
            mask_feature_box = F.interpolate(
                mask_feature_box, size=(align_size,align_size), mode="bilinear", align_corners=False
            )
            boxex_features = torch.cat((boxex_features,mask_feature_box),dim=0)     
        
        zzf = self.aggregator(boxex_features)
        return zzf